import layers
import utils
import dataset

import time
import paddle
from visualdl import LogWriter

#创建模型
vae = layers.VAE()

#设置训练速率策略
#语法参考：https://www.paddlepaddle.org.cn/documentation/docs/zh/api/paddle/optimizer/lr/StepDecay_cn.html
scheduler = paddle.optimizer.lr.StepDecay(learning_rate=0.1, step_size=1, gamma=0.5, verbose=True)
#设置优化器
optim = paddle.optimizer.Adam(learning_rate=\
                              scheduler,\
                              beta1=0.9,\
                              beta2=0.999,\
                              epsilon=1e-08,\
                              parameters=vae.parameters())
    
#损失函数
loss_fn = paddle.nn.MSELoss()
#加载训练数据集
train_loader = paddle.io.DataLoader(dataset.pokemonDataset(dataset_path="E:/Datasets/pokemonJPG"),batch_size=12,shuffle=True)

#开始训练
step = 0
for epoch in range(100000):
    for batch,data in enumerate(train_loader()):
        step += 1
        start = time.time()
        x_data = paddle.to_tensor(data[0],dtype="float32")  #数据
        y_data = paddle.to_tensor(data[1],dtype="float32")  #标签
        predicts = vae(x_data)  #预测结果
        #计算损失
        loss = loss_fn(predicts,y_data)
        #反向传播
        loss.backward()
        #更新参数
        optim.step()
        #梯度清零
        optim.clear_grad()
        #将损失添加给self.losses
        print("\rstep=",step,"loss=",loss.numpy()[0],"cost:",time.time()-start,"s",end="")
        
        if (step+1)%10==0:
            loss = loss.numpy()[-1]
            result = predicts.numpy()
            for i in range(len(result)):
                img = result[i]*255
                img = utils.channel_first_to_last(img)
                utils.save_bgr(img,"result/%d.jpg"%(i))
            with LogWriter(logdir="log/loss_1") as writer:
                writer.add_scalar(tag="loss", step=step, value=loss)

            paddle.save(vae.state_dict(),"parameters/%d_loss_%f.pdparams"%(step,loss))
            paddle.save(optim.state_dict(),"parameters/%d_loss_%f.pdopt"%(step,loss))
    #检测是否减半训练速率
    if step%10000==0:
        scheduler.step()



